//
//  AVLTree.swift
//  数据结构与算法
//
//  Created by Rowling on 2024/2/27.
//

import Foundation
class AVLNode<T> : Node<T> {
    var height = 0
}

class AVLTree<T : Comparable> : BinarySearchTree<T> {
    
    override func afterAdd(node: Node<T>) {
        super.afterAdd(node: node)
        var theNode :  Node<T>? = node
        while (theNode != nil) {
            // 每次添加之后,更新高
            let p = theNode?.parent
            refreshHeight(node: theNode as! AVLNode<T>)
            rebanlance(node: theNode as! AVLNode<T>)
            theNode = p
        }
    }
    
    override func afterRemove(node: Node<T>, replacement: Node<T>?) {
        super.afterAdd(node: node)
        var theNode :  Node<T>? = node
        while (theNode != nil) {
            // 每次添加之后,更新高
            let p = theNode?.parent
            refreshHeight(node: theNode as! AVLNode<T>)
            rebanlance(node: theNode as! AVLNode<T>)
            theNode = p
        }
    }
    
    // 计算平衡因子
    func getFactor(node: Node<T>) -> Int {
        guard let aNode = node as? AVLNode else {return 0}
        let leftF = (aNode.left as? AVLNode)?.height ?? 0
        let rightF = (aNode.right as? AVLNode)?.height ?? 0
        return leftF - rightF
    }
    
    override func createNode(element: T, left: Node<T>? = nil, right: Node<T>? = nil, parent: Node<T>? = nil) -> Node<T> {
        return AVLNode(element: element, left: left, right: right, parent: parent)
    }
    
    func refreshHeight(node: AVLNode<T>) {
        node.height = 1 + max((node.left as? AVLNode)?.height ?? 0, (node.right as? AVLNode)?.height ?? 0)
    }
    
    func rebanlance(node: AVLNode<T>) {
        let factor = getFactor(node: node)
        if abs(factor) <= 1 { return }
        // 不平衡
        if factor > 0 {
            // L
            let subF = getFactor(node: node.left!) // 获取子节点平衡因子
            if subF > 0 {
                // LL
                rightRotation(node: node)
            } else {
                // LR
                leftRotation(node: node.left!)
                rightRotation(node: node)
            }
        } else {
            // R
            let subF = getFactor(node: node.right!) // 获取子节点平衡因子
            if subF > 0 {
                // RL
                rightRotation(node: node.right!)
                leftRotation(node: node)
            } else {
                // RR
                leftRotation(node: node)
            }
        }
    }
    //TODO: 统一旋转操作
//    func rotation(node : Node<T>, left : Bool) {
//        guard let subNode = left ? node.right else {return}
//    }
    func leftRotation(node : Node<T>) {
        guard let right = node.right else {return}
        if node === root {
            root = right
        } else if let p = node.parent {
            if node.isLeft {
                p.left = right
            } else if node.isRight {
                p.right = right
            }
        }
        right.parent = node.parent
        node.parent = right
        node.right = right.left
        right.left?.parent = node
        right.left = node
        refreshHeight(node: node as! AVLNode)
        refreshHeight(node: right as! AVLNode)
    }
    
    
    
    func rightRotation(node : Node<T>) {
        guard let left = node.left else {return}
        if node === root {
            root = left
        } else if let p = node.parent {
            if p.left === node {
                p.left = left
            } else if p.right === node {
                p.right = left
            }
        }
        left.parent = node.parent
        node.parent = left
        node.left = left.right
        left.right?.parent = node
        left.right = node
        refreshHeight(node: node as! AVLNode)
        refreshHeight(node: left as! AVLNode)
    }
    
}
